import torch
import torch.nn as nn
from matplotlib import pyplot as plt

# 准备数据
# 10组数据，每组数据有2个特征
x1 = torch.arange(10).float().unsqueeze(1)
x2 = torch.arange(10, 20).float().unsqueeze(1)
x = torch.concatenate([x1, x2], dim=1)
y = x1 * 0.5 + x2 * 1
# 目前x和y是标准线性，实际上y的值是有一些抖动的，现在来模拟抖动
y += torch.normal(0, 0.1, y.shape)

# 1. 超参数
lr = 0.0001
Epochs = 1000

# 2. 创建模型——输入特征数为2，输出特征数为1
model = nn.Linear(2, 1, bias=False)

# 3. 创建损失函数
loss_fn = nn.MSELoss()

# 4. 创建优化器 (SGD--梯度下降)
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

fig = plt.figure("3D", (12, 6))
ax1 = fig.add_subplot(1, 2, 1, projection="3d")
ax2 = fig.add_subplot(1, 2, 2)
epoch_list = []
loss_list = []
# 5. 循环训练
for epoch in range(Epochs):
    # 前向传播
    y_pre = model(x)
    # 计算损失
    loss = loss_fn(y, y_pre)
    # 反向传播
    loss.backward()
    # 更新参数
    optimizer.step()
    # 清零梯度
    optimizer.zero_grad()

    epoch_list.append(epoch)
    loss_list.append(loss.item())
    # 打印训练信息
    if epoch == 0 or (epoch + 1) % 10 == 0:
        print(f'[{epoch + 1}/{Epochs}], loss: {loss.item():.4f}')
        ax1.clear()
        ax1.scatter(x1, x2, y, c="r")
        ax1.scatter(x1, x2, y_pre.detach(), c="b")

        ax2.clear()
        ax2.plot(epoch_list, loss_list)

        plt.pause(0.1)

# 打印训练得到的参数 w≈0.5, w2≈1
# print(loss_list)
print(f'w1: {model.weight[0][0].item()}, w2: {model.weight[0][1].item()}')

# 6. 学生练习，绘制动态拟合图（循环更新） 和 损失曲线图
plt.show()
